#include "PlaneMeshlight.h"

using namespace OBR;

std::shared_ptr<CudaProgram> PlaneMeshLight::program = nullptr;
CudaBuffer PlaneMeshLight::sbt_data_payload;

PlaneMeshLight::PlaneMeshLight() : AreaLight("shaders/Light/PlaneMeshLight.cu")
{
    // init program opt & desc
    this->program_desc.callables.entryFunctionNameDC = "__direct_callable__rectangle_light";
    this->program_desc.kind = OPTIX_PROGRAM_GROUP_KIND_CALLABLES;
    this->program_desc.flags = OPTIX_PROGRAM_GROUP_FLAGS_NONE;
}

PlaneMeshLight::~PlaneMeshLight()
{
}

void PlaneMeshLight::_buildFromConfig(const Json::Value &config)
{
    ConfigHelpers::loadSpectrum(config["radiance"], &this->radiance);
}

void PlaneMeshLight::uploadVirtualProgramPayload(
    std::vector<std::shared_ptr<VirtualProgramEntry>> instances)
{
    const size_t n_instances = instances.size();
    const size_t data_size = sizeof(PlaneMeshLightData);
    const size_t memory_size = n_instances * data_size;

    // collect data
    PlaneMeshLightData *tmpdata = new PlaneMeshLightData[n_instances];
    for (size_t i = 0; i < n_instances; i++)
    {
        auto l = std::dynamic_pointer_cast<PlaneMeshLight>(instances[i]);
        tmpdata[i].radiance = l->radiance;
        tmpdata[i].l2w = l->local2world;
        tmpdata[i].w2l = l->local2world.inv();
        if (l->mesh != nullptr)
        {
            std::shared_ptr<PlaneMesh> m = std::dynamic_pointer_cast<PlaneMesh>(l->mesh);
            tmpdata[i].corner = m->O;
            tmpdata[i].U = m->U;
            tmpdata[i].V = m->V;
            tmpdata[i].normal_world = normalize(tmpdata[i].w2l.transpose() * m->N);
            tmpdata[i].area = m->area;
        }
    }

    // upload data and free tmp
    PlaneMeshLight::sbt_data_payload.resize(memory_size);
    PlaneMeshLight::sbt_data_payload.upload(tmpdata, memory_size);
    delete[] tmpdata;
}

CUdeviceptr PlaneMeshLight::getPayloadDevicePtr()
{
    return PlaneMeshLight::sbt_data_payload.d_pointer();
}

void PlaneMeshLight::setProgram(std::shared_ptr<CudaProgram> prog)
{
    PlaneMeshLight::program = prog;
}

std::shared_ptr<CudaProgram> PlaneMeshLight::getProgram()
{
    return PlaneMeshLight::program;
}

OptixProgramGroupOptions &PlaneMeshLight::getProgramGroupOptions()
{
    return this->program_opt;
}

OptixProgramGroupDesc &PlaneMeshLight::getProgramGroupDesc()
{
    this->program_desc.callables.moduleDC = this->cuda_module->optix_module;
    return this->program_desc;
}

void PlaneMeshLight::updateSbtData()
{
    assert(this->sbt_data_payload.d_ptr != nullptr);
    PlaneMeshLightData *d_loc =
        (PlaneMeshLightData *)this->sbt_data_payload.d_ptr + this->program_instance_id;
    PlaneMeshLightData data;
    std::shared_ptr<PlaneMesh> m = std::dynamic_pointer_cast<PlaneMesh>(this->mesh);
    data.corner = m->O;
    data.U = m->U;
    data.V = m->V;
    data.normal_world = normalize(data.w2l.transpose() * m->N);
    data.area = m->area;
    CUDA_CHECK(cudaMemcpy(d_loc, &data, sizeof(PlaneMeshLightData), cudaMemcpyHostToDevice));
}

void PlaneMeshLight::linkMesh(std::shared_ptr<Mesh> _mesh)
{
    try
    {
        std::shared_ptr<PlaneMesh> _m = std::dynamic_pointer_cast<PlaneMesh>(_mesh);
        this->mesh = _mesh;
    }
    catch (const std::exception &e)
    {
        std::cerr << e.what() << '\n';
        throw e;
    }
}